import platform
import pandas as pd
import torch
from trade_agent import TradeAgent

num_envs = 20

model_save_path = '/root/autodl-fs/OKX-BTC-USDT-SWAP-1s.zip'
data_path = "/root/autodl-fs/OKX-BTC-USDT-SWAP-1s-features.csv"
system_name = platform.system()
if system_name == "Windows":
    num_envs = 1
    model_save_path = './OKX-BTC-USDT-SWAP-1s_1.zip'
    data_path = "./OKX-BTC-USDT-SWAP-1s-features.csv"

TRAIN_RATIO = 0.8
step_multiples = 4

model_kwargs = dict(
    learning_rate = 1e-5,
    buffer_size = 1_000_000,
    batch_size = 1024,
    train_freq = 64,
    gradient_steps = 8,
    target_update_interval = 8000,
    exploration_fraction = 0.1,
    exploration_initial_eps = 1.0,
    exploration_final_eps = 0.02,
    gamma = 1, #0.99997, # 一万步后 0.7  八万步后 0.07
    tau = 0.005,

    policy_kwargs = dict(
        net_arch=[256, 128],
        activation_fn=torch.nn.SiLU,
        features_extractor_kwargs=dict(
            seq_len=90,
            d_model=128,
            nhead=4,
            num_layers=3,
        )
    ),
)

custom_objects = {
    "lr_schedule": lambda _: 1e-5,
    "exploration_schedule": lambda _: 0.02,
}

all_indicator_names = [
    "macd", "boll_ub", "boll_lb", "rsi_30", "cci_30", "dx_30",
    "close_30_sma", "close_60_sma", "ema_12", "ema_26",
    "momentum_10", "roc_10", "willr_14", "stoch_k_14", "stoch_d_14",
    "atr_14", "trange", "obv", "vwap"
]

def main():
    base_data = pd.read_csv(data_path)
    print(f"原始数据长度: {len(base_data)}")

    train_df = base_data.iloc[:int(len(base_data) * TRAIN_RATIO)]

    agent = TradeAgent(
        df=train_df,
        tech_indicator_list=all_indicator_names,
        model_kwargs=model_kwargs,
        num_envs=num_envs
    )
    agent.train_model(
        model_save_path=model_save_path,
        step_multiples=step_multiples,
        custom_objects=custom_objects,
    )

if __name__ == "__main__":
    main()